import argparse
import os
import sys
from pathlib import Path

import torch
import numpy as np
import torchvision.transforms as transforms

from torch.optim import AdamW
from mmengine.optim import AmpOptimWrapper
from mmengine.runner import Runner

from camvid.metric import IoU
from camvid.hook import SegVisHook
from camvid.dataset import CamVid
from camvid.model import MMDeeplabV3

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative


def main(opt):
    norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    transform = transforms.Compose(
        [transforms.ToTensor(), transforms.Normalize(**norm_cfg)]
    )

    target_transform = transforms.Lambda(
        lambda x: torch.tensor(np.array(x), dtype=torch.long)
    )

    train_set = CamVid(
        opt.data_dir,
        img_folder="train",
        mask_folder="train_labels",
        transform=transform,
        target_transform=target_transform,
    )

    valid_set = CamVid(
        opt.data_dir,
        img_folder="val",
        mask_folder="val_labels",
        transform=transform,
        target_transform=target_transform,
    )

    train_dataloader = dict(
        batch_size=3,
        dataset=train_set,
        sampler=dict(type="DefaultSampler", shuffle=True),
        collate_fn=dict(type="default_collate"),
    )

    val_dataloader = dict(
        batch_size=3,
        dataset=valid_set,
        sampler=dict(type="DefaultSampler", shuffle=False),
        collate_fn=dict(type="default_collate"),
    )

    num_classes = 32  # Modify to actual number of categories.

    runner = Runner(
        model=MMDeeplabV3(num_classes),
        work_dir=opt.work_dir,
        train_dataloader=train_dataloader,
        optim_wrapper=dict(type=AmpOptimWrapper, optimizer=dict(type=AdamW, lr=2e-4)),
        train_cfg=dict(by_epoch=True, max_epochs=opt.epochs, val_interval=10),
        val_dataloader=val_dataloader,
        val_cfg=dict(),
        val_evaluator=dict(type=IoU),
        custom_hooks=[SegVisHook(opt.data_dir)],
        default_hooks=dict(checkpoint=dict(type="CheckpointHook", interval=1)),
    )
    runner.train()


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--data_dir", default=ROOT / "data/CamVid", help="the data directory"
    )
    parser.add_argument(
        "--work_dir", default=ROOT / "runs/camvid", help="the working directory"
    )
    parser.add_argument("--epochs", default=10, type=int, help="epochs to train")
    opt = parser.parse_args()
    return opt


if __name__ == "__main__":
    opt = parse_opt()
    main(opt)
